/*
 * TotalEquilibriumProblem.cpp
 *
 *  Created on: 4 Sep 2011
 *      Author: Allan
 */

#include "TotalEquilibriumProblem.h"

TotalEquilibriumProblem::TotalEquilibriumProblem()
{}

TotalEquilibriumProblem::TotalEquilibriumProblem(const Multiphase& multiphase, const CanonicalReactionSystem& reactions) :
multiphase(multiphase), reactions(reactions)
{
	// The primary (jSpecies) and equilibrium (jSpecies) species
	const vector<string> jSpecies = reactions.GetPrimarySpecies();
	const vector<string> eSpecies = reactions.GetSecondarySpecies();
	
	// Set the number of primary (Nj) and equilibrium (Ne) species
	Nj = jSpecies.size();
	Ne = eSpecies.size();
	
	// Set the indexes of the primary (jIndexes) and equilibrium (eIndexes) species
	jIndexes = multiphase[jSpecies];
	eIndexes = multiphase[eSpecies];
	
	// Set the stoichiometric matrix of the system of reactions
	vej = reactions.GetCanonicalStoichiometricMatrix();
	
	// Allocate all the auxiliary vector/matrix variables with zero entries
	Ke = VectorXd::Zero(Ne);
	Qe = VectorXd::Zero(Ne);
	ne = VectorXd::Zero(Ne);
}

TotalEquilibriumProblem::~TotalEquilibriumProblem()
{}

void TotalEquilibriumProblem::SetParameters(const VectorXd& uj, double T, double P, const VectorXd& n)
{
	this->T  = T;
	this->P  = P;
	this->n  = n;
	this->uj = uj;
	
	// Initialise the equilibrium constants of the equilibrium reactions (Ke)
	Ke = reactions.CanonicalEquilibriumConstants(T, P);
	
	// Initialise the vector (ne) with the values from the initial guess (n)
	for(unsigned e = 0; e < Ne; ++e) ne[e] = n[eIndexes[e]];
}

void TotalEquilibriumProblem::Function(const VectorXd& nj, VectorXd& F)
{
	// Update the auxiliary data members of the equilibrium problem
	UpdateAuxiliaryData(nj);
	
	// Compute the residual vector (F)
	F = nj + vej.transpose() * ne - uj;
}

void TotalEquilibriumProblem::Jacobian(const VectorXd& nj, MatrixXd& J)
{
	// Compute the jacobian matrix (J)
	for(unsigned i = 0; i < Nj; ++i) for(unsigned j = 0; j < Nj; ++j)
	{
		double sum = 0.0; 
		
		for(unsigned e = 0; e < Ne; ++e) 
			sum += vej(e, i) * vej(e, j) * ne[e];
		
		J(i, j) = sum/nj[j];
	}
	
	J += MatrixXd::Identity(Nj, Nj);
	
	//const VectorXd inv_nj = nj.inverse();
	//J  = vej.transpose() * ne.asDiagonal() * vej * inv_nj.asDiagonal();
	//J += MatrixXd::Identity(Nj, Nj);
}

void TotalEquilibriumProblem::UpdateAuxiliaryData(const VectorXd& nj)
{
	// Check if all the entries in (nj) are non-negative
	if(nj.minCoeff() > 0)
	{
		// Update the nj part of the vector (n) with the values of the vector (nj)
		for(unsigned j = 0; j < Nj; ++j) n[jIndexes[j]] = nj[j];
		
		// Update the activity vector (a) with temperature (T), pressure (P) and the mole vector (n)
		a = multiphase.Activities(T, P, n);
	
		// Update the reaction quotients of the equilibrium reactions (Qe)
		Qe = reactions.CanonicalReactionQuotients(a);
	
		// Update the moles of the equilibrium species (ne)
		ne = ne.array() * Qe.array() / Ke.array();
	
		// Update the ne part of the vector (n)
		for(unsigned e = 0; e < Ne; ++e) n[eIndexes[e]] = ne[e];
	}
}
